{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 迭代重写"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax as rx\n",
    "from tvm import relay, tir\n",
    "from tvm.relax.analysis import get_var2val\n",
    "import tvm.testing\n",
    "from tvm.relax.dpl import *\n",
    "\n",
    "same_shape_func_type = tvm.testing.parameter(\n",
    "    \"same_static_shape\",\n",
    "    \"same_dynamic_shape\",\n",
    "    \"different_static_shape\",\n",
    "    \"different_dynamic_shape\",\n",
    ")\n",
    "def test_iterative_rewrite_without_trivial_binding():\n",
    "    \"\"\"Avoid introducing common sub-expressions\n",
    "\n",
    "    Pattern replacement may produce the same intermediate, which\n",
    "    should appear only once in the final result.\n",
    "    \"\"\"\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before(x: R.Tensor((1024,))):\n",
    "        with R.dataflow():\n",
    "            a = R.strided_slice(x, [0], [0], [512], [1])\n",
    "            b = R.strided_slice(x, [0], [512], [1024], [1])\n",
    "            c = R.add(a, b)\n",
    "            R.output(c)\n",
    "        return c\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected(x: R.Tensor((1024,))):\n",
    "        with R.dataflow():\n",
    "            x_split = R.split(x, 2)\n",
    "            a = x_split[0]\n",
    "            b = x_split[1]\n",
    "            c = R.add(a, b)\n",
    "            R.output(c)\n",
    "        return c\n",
    "\n",
    "    pattern_arg = wildcard()\n",
    "    pattern_axes = wildcard()\n",
    "    pattern_begin = wildcard()\n",
    "    pattern_end = wildcard()\n",
    "    pattern_strides = wildcard()\n",
    "    pattern = is_op(\"relax.strided_slice\")(\n",
    "        pattern_arg, pattern_axes, pattern_begin, pattern_end, pattern_strides\n",
    "    )\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        arg = matches[pattern_arg]\n",
    "        axes = matches[pattern_axes]\n",
    "        begin = matches[pattern_begin]\n",
    "        end = matches[pattern_end]\n",
    "        strides = matches[pattern_strides]\n",
    "        strided_slice = matches[pattern]\n",
    "\n",
    "        if arg.struct_info.shape is None:\n",
    "            return expr\n",
    "\n",
    "        if len(axes) != 1:\n",
    "            return expr\n",
    "\n",
    "        axis = axes[0].value\n",
    "        begin = begin[0].value\n",
    "        end = end[0].value\n",
    "        stride = strides[0].value\n",
    "\n",
    "        if stride != 1:\n",
    "            return expr\n",
    "\n",
    "        size = arg.struct_info.shape[0]\n",
    "        if (\n",
    "            isinstance(size, tir.IntImm)\n",
    "            and isinstance(begin, tir.IntImm)\n",
    "            and isinstance(end, tir.IntImm)\n",
    "        ):\n",
    "            size = size.value\n",
    "            begin = begin.value\n",
    "            end = end.value\n",
    "        else:\n",
    "            return expr\n",
    "\n",
    "        gcd = functools.reduce(math.gcd, [begin, end, size])\n",
    "        if (end - begin) // gcd == 1:\n",
    "            return rx.op.split(arg, size // gcd)[begin // gcd]\n",
    "\n",
    "        return expr\n",
    "\n",
    "    after = rewrite_call(pattern, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(after, expected)\n",
    "\n",
    "\n",
    "def test_iterative_rewrite_with_removed_intermediates():\n",
    "    \"\"\"Pattern replacement may require canonicalization\n",
    "\n",
    "    A pattern may replace a tuple returned by a function with a tuple\n",
    "    whose contents are known by Relax.  In that case, canonicalization\n",
    "    is required to unwrap the TupleGetItem instances into the known\n",
    "    contents.\n",
    "\n",
    "    This test case shows the intermediate results produced in the\n",
    "    process of pattern-matching.\n",
    "    \"\"\"\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "        with R.dataflow():\n",
    "            c = R.concat([a, b])\n",
    "            d = R.split(c, 2)\n",
    "            e = d[0]\n",
    "            f = d[1]\n",
    "            g = R.add(a, e)\n",
    "            h = R.add(f, g)\n",
    "            R.output(h)\n",
    "        return h\n",
    "\n",
    "    # First pattern rewrite.  The concat/rewrite can be unwrapped, so\n",
    "    # `d` is rewritten from `R.split(c, 2)` into `(a, b)`.\n",
    "    #\n",
    "    # @R.function(private=True)\n",
    "    # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "    #     with R.dataflow():\n",
    "    #         c = R.concat([a, b])\n",
    "    #         d = (a,b)\n",
    "    #         e = d[0]\n",
    "    #         f = d[1]\n",
    "    #         g = R.add(a, e)\n",
    "    #         h = R.add(f, g)\n",
    "    #         R.output(h)\n",
    "\n",
    "    # Canonicalization step.  Because `d` is known to be `(a,b)`,\n",
    "    # canonicalization can rewrite `d[0]` into `a` and `d[1]` into\n",
    "    # `b`.\n",
    "    #\n",
    "    # @R.function(private=True)\n",
    "    # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "    #     with R.dataflow():\n",
    "    #         c = R.concat([a, b])\n",
    "    #         d = (a,b)\n",
    "    #         e = a\n",
    "    #         f = b\n",
    "    #         g = R.add(a, a)\n",
    "    #         h = R.add(b, g)\n",
    "    #         R.output(h)\n",
    "\n",
    "    # Dead-code-elimination step.  This technically isn't required\n",
    "    # until the pattern matching has converged, but performing it now\n",
    "    # prevents testing for matches on dead code.\n",
    "    #\n",
    "    # @R.function(private=True)\n",
    "    # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "    #     with R.dataflow():\n",
    "    #         g = R.add(a, a)\n",
    "    #         h = R.add(b, g)\n",
    "    #         R.output(h)\n",
    "\n",
    "    # Second pattern-matching step.  Now, the `R.add(a,a)` can match\n",
    "    # the other option in our pattern, and be rewritten as\n",
    "    # `R.multiply(a,R.const(2))`.\n",
    "    #\n",
    "    # @R.function(private=True)\n",
    "    # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "    #     with R.dataflow():\n",
    "    #         g = R.multiply(a, R.const(2))\n",
    "    #         h = R.add(b, g)\n",
    "    #         R.output(h)\n",
    "\n",
    "    # Canonicalization and dead-code-elimination are applied again,\n",
    "    # but have no effect this time.\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected(a: R.Tensor((1024,)), b: R.Tensor((1024,))):\n",
    "        with R.dataflow():\n",
    "            g = R.multiply(a, R.const(2))\n",
    "            h = R.add(b, g)\n",
    "            R.output(h)\n",
    "        return h\n",
    "\n",
    "    pat_args = wildcard()\n",
    "\n",
    "    op_concat = is_op(\"relax.concat\")\n",
    "    pat_concat = op_concat(pat_args).has_attr({\"axis\": 0})\n",
    "\n",
    "    op_split = is_op(\"relax.split\")\n",
    "    pat_split = op_split(pat_concat).has_attr({\"axis\": 0, \"indices_or_sections\": T.int64(2)})\n",
    "\n",
    "    pat_unwrap_concat_split = pat_split\n",
    "\n",
    "    pat_arg = wildcard()\n",
    "    op_add = is_op(\"relax.add\")\n",
    "    pat_add_self = op_add(pat_arg, pat_arg)\n",
    "\n",
    "    pattern = pat_unwrap_concat_split | pat_add_self\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        if pat_unwrap_concat_split in matches:\n",
    "            args = matches[pat_args]\n",
    "\n",
    "            if len(args) == 2 and tvm.ir.structural_equal(args[0].struct_info, args[1].struct_info):\n",
    "                return args\n",
    "\n",
    "        elif pat_add_self in matches:\n",
    "            arg = matches[pat_arg]\n",
    "            return arg * rx.const(2)\n",
    "\n",
    "        return expr\n",
    "\n",
    "    after = rewrite_call(pattern, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
